# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import pathlib
import platform
import math
import numpy as np

import dateutil.parser
from absl import app, flags

from coin.base.datetime_util import iterate_date
from coin.experimental.xguo.convert_symbol_to_product import convert_symbol_to_product
from coin.experimental.xguo.fruit.util.feed_util import (
    AutoInc,
    load_products_feed_as_json,
    get_moving_average,
    get_moving_diff,
    get_moving_return,
    get_mid_series,
    get_trade_qty_series,
    to_mid_night,
)

#matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def plot_result1(prices):
  plt.figure(figsize=(20, 10))
  ref = None
  for product, price in prices.items():
    if ref is None:
      ref = price.median()
      print(product, ', ref: ', ref)
      plt.plot(price.index, price, label=product.full_symbol)
    else:
      fac = ref / price.median()
      print(product, ', fac: ', fac)
      plt.plot(price.index, price, label=product.full_symbol)

  plt.legend()
  plt.show()


def plot_result2(prices):
  plt.figure(figsize=(20, 10))
  freq = datetime.timedelta(seconds=1)
  res = {}
  for product, price in prices.items():
    price = price.asfreq(freq, method='backfill')
    if res.get(product.base.currency):
      res[product.base.currency].append(price)
    else:
      res[product.base.currency] = [price]

  for base, prices in res.items():
    x = prices[0] / prices[1]
    plt.plot(get_moving_average(x, window=300), label=base)
  plt.legend()
  plt.show()


def plot_result(
    home_dir,
    trading_dates,
    product,
    mid_price_series,
    trade_qty_series,
):
  fig_root_dir = home_dir.joinpath('feed_research')
  trading_date = trading_dates[0].strftime('%Y%m%d')
  plt.figure(figsize=(20, 30))
  idx = AutoInc(start=1)
  num_rows = 7

  freq = datetime.timedelta(seconds=1)
  freq1 = datetime.timedelta(seconds=1)
  s10 = datetime.timedelta(seconds=10)
  s300 = datetime.timedelta(seconds=300)
  s600 = datetime.timedelta(seconds=600)
  s1200 = datetime.timedelta(seconds=1200)
  s2400 = datetime.timedelta(seconds=2400)
  s3600 = datetime.timedelta(seconds=3600)

  plt.subplot(num_rows, 1, idx())
  plt.plot(mid_price_series.index, mid_price_series, label='orig')
  x1 = get_moving_average(mid_price_series, s1200)
  x2 = get_moving_average(mid_price_series, s2400)
  x3 = get_moving_average(mid_price_series, s3600)

  plt.plot(x1.index, x1, label='1200')
  plt.plot(x2.index, x2, label='2400')
  plt.plot(x2.index, x3, label='3600')
  ax = plt.twinx()
  g1 = trade_qty_series.resample(freq).sum()
  ax.plot(g1.index, g1)

  plt.subplot(num_rows, 1, idx())
  diff_prices = get_moving_return(mid_price_series, s300)
  diff_prices = diff_prices.asfreq(freq1, method='backfill')
  N = int(2**math.ceil(math.log2(len(diff_prices))))
  y = np.abs(np.fft.fft(diff_prices[1:], N))
  y = np.fft.fftshift(y)
  fs = 1 / 300
  x = (np.array(range(N)) - N / 2) / N * fs
  plt.plot(x, y)

  plt.subplot(num_rows, 1, idx())
  prices = x1 - x2
  plt.plot(prices.index, prices, label='1200-2400')
  prices = x1 - x3
  plt.plot(prices.index, prices, label='1200-3600')
  plt.legend()

  plt.subplot(num_rows, 1, idx())
  ser = get_moving_return(mid_price_series, s300)
  plt.plot(ser.index, ser, label='1200', marker='.')
  ma = get_moving_average(ser, s300)
  plt.plot(ma.index, ma, label='MA1200', marker='.')
  ma = get_moving_average(ser, s600)
  plt.plot(ma.index, ma, label='MA2400', marker='.')
  y = 0.002
  plt.plot([ma.index[0], ma.index[-1]], [y, y])
  plt.plot([ma.index[0], ma.index[-1]], [-y, -y])
  plt.legend()

  plt.subplot(num_rows, 1, idx())
  x1 = get_moving_average(trade_qty_series, s1200)
  x2 = get_moving_average(trade_qty_series, s2400)
  x3 = get_moving_average(trade_qty_series, s3600)
  qty = x1 - x2
  plt.plot(qty.index, qty)
  qty = x1 - x3
  plt.title('trade_qty')
  plt.plot(qty.index, qty)

  plt.subplot(num_rows, 1, idx())
  plt.plot(trade_qty_series.index, trade_qty_series)

  plt.subplot(num_rows, 1, idx())
  plt.plot(trade_qty_series.index, trade_qty_series.cumsum())

  if 'hive' in platform.node():
    fig_parent = fig_root_dir.joinpath(trading_date)
    fig_parent.mkdir(parents=True, exist_ok=True)
    fig_file = fig_parent.joinpath(product.full_symbol + '.png')
    plt.savefig(fig_file.open(mode='wb'))
    plt.close()
  else:
    plt.show()


def main(_):
  if FLAGS.home_dir is None:
    if 'hive' in platform.node():
      home_dir = pathlib.Path('/remote/iosg/home/xguo')
    else:
      home_dir = pathlib.Path.home()
  else:
    home_dir = pathlib.Path(FLAGS.home_dir)

  root_dir = home_dir.joinpath('converted_book')

  products = []
  for symbol in FLAGS.symbol.split(','):
    product = convert_symbol_to_product(symbol)
    products.append(product)

  if '-' in FLAGS.trading_date:
    start_date, end_date = FLAGS.trading_date.split('-')
    start_date = dateutil.parser.parse(start_date)
    end_date = dateutil.parser.parse(end_date)
    trading_dates = list(iterate_date(start_date, end_date))
  else:
    trading_dates = [dateutil.parser.parse(FLAGS.trading_date)]

  feeds = load_products_feed_as_json(root_dir, trading_dates, products)
  prices = {}
  trade_qty = {}
  for product, feed_list in feeds.items():
    prices[product] = get_mid_series(feed_list)
    trade_qty[product] = get_trade_qty_series(feed_list)

  plot_result2(prices)
  """
  product = products[0]
  plot_result(
    home_dir,
    trading_dates,
    product,
    prices[product],
    trade_qty[product],
  )
  """


if __name__ == '__main__':
  flags.DEFINE_string(
      'symbol',
      None,
      'Specify product full symbol.',
  )

  flags.DEFINE_string('trading_date', None, 'Trading date.')

  flags.DEFINE_string('home_dir', None, 'Home directory.')
  app.run(main)
